#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
分段立方插值（μ‑z 数列法）实现 4 倍图像放大
"""

from PIL import Image
import numpy as np

# -------------------------------------------------
# 1️⃣  μ‑z 立方插值核（单维）
# -------------------------------------------------
def mu_z_kernel(t: float) -> np.ndarray:
    """
    计算 μ‑z 立方插值的 4 个权重。
    t : 归一化位置，0 <= t < 1，表示在相邻两个原始采样点之间的相对距离。
    返回: 长度为 4 的 ndarray，分别对应左‑2、左‑1、右‑0、右‑1 四个基点的权重。
    """
    # 立方卷积的标准系数（Catmull‑Rom / B‑spline 等可自行替换）
    # 这里使用 Catmull‑Rom（常用于图像放大），对应 μ‑z 数列的实现。
    t2 = t * t
    t3 = t2 * t

    w0 = -0.5 * t3 + t2 - 0.5 * t          # left‑2
    w1 =  1.5 * t3 - 2.5 * t2 + 1.0       # left‑1
    w2 = -1.5 * t3 + 2.0 * t2 + 0.5 * t    # right‑0
    w3 =  0.5 * t3 - 0.5 * t2              # right‑1
    return np.array([w0, w1, w2, w3], dtype=np.float64)


def cubic_interp_1d(pixels: np.ndarray, scale: int = 4) -> np.ndarray:
    """
    对一维像素序列做 μ‑z 立方插值（每段使用 4×4 像素块）。
    pixels : 长度为 N 的 1‑D ndarray，支持标量（灰度）或向量（RGB）。
    scale  : 放大倍数，固定为 4。
    返回   : 长度为 (N‑1)*scale+1 的 ndarray，已完成插值。
    """
    N = len(pixels)

    # 为了在首尾也能取到 4‑point 块，向两端各复制两个像素
    extended = np.concatenate(
        (pixels[0:2][::-1], pixels, pixels[-2:][::-1]), axis=0
    )  # shape (N+4, …)

    out_len = (N - 1) * scale + 1
    if pixels.ndim == 2:                     # RGB
        out = np.empty((out_len, pixels.shape[1]), dtype=np.float64)
    else:                                    # 灰度
        out = np.empty(out_len, dtype=np.float64)

    # 把原始采样点拷贝到对应位置
    out[::scale] = pixels

    # 对每两个相邻原始点之间插值
    for i in range(N - 1):
        # 取四个基点在 extended 中的索引：i, i+1, i+2, i+3
        base = i
        block = extended[base:base + 4]      # shape (4, …)

        for k in range(1, scale):
            t = k / scale                     # 归一化位置 0<t<1
            w = mu_z_kernel(t)                # (4,)

            # 加权求和
            out[i * scale + k] = np.tensordot(w, block, axes=1)

    return out


# -------------------------------------------------
# 2️⃣  主放大函数（先横向后纵向）
# -------------------------------------------------
def upscale_cubic_4x(img: Image.Image) -> Image.Image:
    """
    将输入图像放大四倍，使用 μ‑z 立方插值（4×4 像素块）。
    参数
    ----
    img : PIL.Image（已转为 RGB 或 L）
    返回
    ----
    out_img : 放大后的 PIL.Image（尺寸 = 4·原尺寸）
    """
    src_w, src_h = img.size
    src_arr = np.asarray(img, dtype=np.float64)   # (h, w, c) 或 (h, w)

    # ---------- 水平插值 ----------
    if src_arr.ndim == 3:          # RGB
        rows = [cubic_interp_1d(row, 4) for row in src_arr]   # (w*4-3, 3)
        horiz = np.stack(rows, axis=0)                       # (h, w*4-3, 3)
    else:                           # 灰度
        rows = [cubic_interp_1d(row, 4) for row in src_arr]
        horiz = np.stack(rows, axis=0)                       # (h, w*4-3)

    # ---------- 纵向插值 ----------
    if horiz.ndim == 3:
        trans = np.transpose(horiz, (1, 0, 2))                # (w', h, c)
        cols = [cubic_interp_1d(col, 4) for col in trans]    # (h*4-3, 3)
        vert = np.stack(cols, axis=0)                         # (w'*4-3, h*4-3, 3)
        final = np.transpose(vert, (1, 0, 2))                 # (h*4-3, w*4-3, 3)
    else:
        trans = np.transpose(horiz)                           # (w', h)
        cols = [cubic_interp_1d(col, 4) for col in trans]
        vert = np.stack(cols, axis=0)                         # (w'*4-3, h*4-3)
        final = np.transpose(vert)                           # (h*4-3, w*4-3)

    # ---------- 补齐到完整 4 倍尺寸 ----------
    dst_h, dst_w = src_h * 4, src_w * 4
    pad_h = dst_h - final.shape[0]
    pad_w = dst_w - final.shape[1]

    if pad_h > 0:
        final = np.vstack([final, np.repeat(final[-1:, :], pad_h, axis=0)])
    if pad_w > 0:
        final = np.hstack([final, np.repeat(final[:, -1:], pad_w, axis=1)])

    # ---------- 转回 Pillow 图像 ----------
    final = np.clip(final, 0, 255).astype(np.uint8)
    out_img = Image.fromarray(final, mode=img.mode)
    return out_img


# -------------------------------------------------
# 3️⃣  示例入口
# -------------------------------------------------
if __name__ == "__main__":
    input_path  = "lena_128.png"          # 128×128 原图
    output_path = "lena_512_cubic.png"

    img = Image.open(input_path).convert("RGB")
    upscaled = upscale_cubic_4x(img)
    upscaled.save(output_path)
    print(f"已完成 4 倍 μ‑z 立方插值放大，保存为 {output_path}")
